import requests
import json
# 本代码输入 4 个函数，分别是数值的加、减、比较大小、字符串中重复字母计数四个函数 来演示通过function calling来解决大语言模型在tokens 预测不擅长的领域的执行问题。
tools = [
    {
        'type': 'function',
        'function': {
            'name': 'add',
            'description': 'Compute the sum of two numbers',
            'parameters': {
                'type': 'object',
                'properties': {
                    'a': {
                        'type': 'int',
                        'description': 'A number',
                    },
                    'b': {
                        'type': 'int',
                        'description': 'A number',
                    },
                },
                'required': ['a', 'b'],
            },
        }
    },
    {
        'type': 'function',
        'function': {
            'name': 'mul',
            'description': 'Calculate the product of two numbers',
            'parameters': {
                'type': 'object',
                'properties': {
                    'a': {
                        'type': 'int',
                        'description': 'A number',
                    },
                    'b': {
                        'type': 'int',
                        'description': 'A number',
                    },
                },
                'required': ['a', 'b'],
            },
        }
    },
    {
        'type': 'function',
        'function': {
            'name': 'count_letter_in_string',
            'description': 'Count letter number in a string',
            'parameters': {
                'type': 'object',
                'properties': {
                    'a': {
                        'type': 'str',
                        'description': 'source string',
                    },
                    'b': {
                        'type': 'str',
                        'description': 'letter',
                    },
                },
                'required': ['a', 'b'],
            },
        }
    },
    {
        'type': 'function',
        'function': {
            'name': 'compare',
            'description': 'Compare two number, which one is bigger',
            'parameters': {
                'type': 'object',
                'properties': {
                    'a': {
                        'type': 'float',
                        'description': 'A number',
                    },
                    'b': {
                        'type': 'float',
                        'description': 'A number',
                    },
                },
                'required': ['a', 'b'],
            },
        }
    }
]
with open("config.json") as f:
    API_KEY = json.load(f)["API_KEY"]
headers = {
    "Authorization": "Bearer sk-" + API_KEY,
    "Content-Type": "application/json"
}
url = "https://api.siliconflow.cn/v1/chat/completions"
def create(messages):
    payload = {
        "model": "Qwen/Qwen2.5-7B-Instruct",
        "messages": messages,
        "stream": False,
        "response_format": {"type": "text"},
        "tools" : tools
    }
    response = requests.request("POST", url, stream=False, json=payload, headers=headers)
    # print(response.text)
    json_data = json.loads(response.text)
    return json_data
    
def add(a: float, b: float):
    return a + b

def mul(a: float, b: float):
        return a * b

def compare(a: float, b: float):
    if a > b:
        return f'{a} is greater than {b}'
    elif a < b:
        return f'{b} is greater than {a}'
    else:
        return f'{a} is equal to {b}'

def count_letter_in_string(a: str, b: str):
    string = a.lower()
    letter = b.lower()

    count = string.count(letter)
    return (f"The letter '{letter}' appears {count} times in the string.")




def function_call_playground(prompt):
    messages = [{ 'role': 'user', 'content': prompt }]
    response = create(messages)
    # print(response)
    # print(response)
    func1_name = response["choices"][0]["message"]["tool_calls"][0]["function"]["name"]
    func1_args = response["choices"][0]["message"]["tool_calls"][0]["function"]["arguments"]
    func1_out = eval(f'{func1_name}(**{func1_args})')
    # print(func1_out)

    messages.append(response["choices"][0]["message"])
    messages.append({
        'role': 'tool',
        'content': f'{func1_out}',
        'tool_call_id': response["choices"][0]["message"]["tool_calls"][0]["id"]
    })
        # print(messages)
    response = create(messages)

    return response["choices"][0]["message"]["content"]

prompts = [
    "用中文回答：strawberry中有多少个r?",
    "用中文回答：9.11和9.9，哪个小?"
]

for prompt in prompts:
    print(function_call_playground(prompt))